import cv2
import numpy as np


def extract_paper(image_path):
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        print(f"无法加载图像: {image_path}")
        return None

    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 使用高斯模糊减少噪声
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 使用Canny边缘检测
    edges = cv2.Canny(blurred, 50, 150)

    # 形态学操作：先膨胀后腐蚀，以去除小噪声并连接断开的部分
    kernel = np.ones((5, 5), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)
    edges = cv2.erode(edges, kernel, iterations=1)

    # 使用霍夫变换检测直线
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)

    if lines is not None:
        # 找到四个角点
        corners = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            corners.append([x1, y1])
            corners.append([x2, y2])

        # 使用最小面积矩形拟合四个角点
        corners = np.array(corners)
        rect = cv2.minAreaRect(corners)
        box = cv2.boxPoints(rect)
        box = np.int32(box)  # 修改这里

        # 得到感兴趣区域（ROI）
        width = int(rect[1][0])
        height = int(rect[1][1])
        src_pts = box.astype("float32")
        dst_pts = np.array([[0, height - 1],
                            [0, 0],
                            [width - 1, 0],
                            [width - 1, height - 1]], dtype="float32")
        M = cv2.getPerspectiveTransform(src_pts, dst_pts)
        warped = cv2.warpPerspective(img, M, (width, height))

        return warped

    return None


if __name__ == "__main__":
    image_path = '/home/weiqiangren/Python/project/py-basic-exercises/图片抠图调整方向/企业微信截图_17388262872342.png'  # 替换为你的图像路径
    extracted_paper = extract_paper(image_path)

    if extracted_paper is not None:
        output_path = 'extracted_paper.png'
        cv2.imwrite(output_path, extracted_paper)  # 保存结果图像
        print(f"矩形白纸已成功提取并保存为 '{output_path}'")

        # 如果你有GUI支持并且想要显示图像，可以取消以下注释：
        # cv2.imshow('Extracted Paper', extracted_paper)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
    else:
        print("未找到矩形白纸.")